"""
python fujian2_lstm_batch_predict.py
"""

import os
import json
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from datetime import datetime, timedelta
from tqdm import tqdm  # Import tqdm for progress bar

# LSTM Dataset class
class TimeSeriesDataset(Dataset):
    def __init__(self, data, sequence_length=10):
        self.data = data
        self.sequence_length = sequence_length

    def __len__(self):
        return len(self.data) - self.sequence_length

    def __getitem__(self, index):
        return (self.data[index:index + self.sequence_length],
                self.data[index + self.sequence_length])

# Define LSTM model
class LSTMModel(nn.Module):
    def __init__(self, input_size=1, hidden_size=64, num_layers=1, output_size=1):
        super(LSTMModel, self).__init__()
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        lstm_out, _ = self.lstm(x)
        return self.fc(lstm_out[:, -1, :])

# Function to load and preprocess data
def load_data(input_dir):
    all_data = {}
    for file in os.listdir(input_dir):
        if file.endswith('.json'):
            with open(os.path.join(input_dir, file), 'r') as f:
                data = json.load(f)
                df = pd.DataFrame(data)
                df['ds'] = pd.to_datetime(df['date'])
                df.set_index('ds', inplace=True)
                df = df.asfreq('D', fill_value=0)  # Fill missing dates with 0
                all_data[file] = df['sales'].values
    return all_data

# Function to train LSTM and make predictions
def train_and_predict(data, output_json_path, graph_output_path):
    sequence_length = 10
    all_sequences = []

    for category, sales in data.items():
        dataset = TimeSeriesDataset(torch.FloatTensor(sales).view(-1, 1), sequence_length)
        all_sequences.append(dataset)

    # Create a DataLoader with combined data
    combined_data = torch.cat([torch.FloatTensor(sales).view(-1, 1) for sales in data.values()], dim=0)
    dataset = TimeSeriesDataset(combined_data, sequence_length)
    data_loader = DataLoader(dataset, batch_size=32, shuffle=True)

    model = LSTMModel()
    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    # Training the model
    for epoch in range(20):
        for X, y in data_loader:
            optimizer.zero_grad()
            output = model(X.view(-1, sequence_length, 1))
            loss = criterion(output.view(-1), y.view(-1))
            loss.backward()
            optimizer.step()

    # Make predictions for each category
    predictions = {}
    future_dates = pd.date_range(start='2023-07-01', periods=92, freq='D')

    # Use tqdm to show progress bar for categories
    for category in tqdm(data.keys(), desc='Processing Categories', unit='category'):
        sales = data[category]
        last_sequence = torch.FloatTensor(sales[-sequence_length:]).view(1, sequence_length, 1)

        category_predictions = []
        for _ in range(92):
            with torch.no_grad():
                pred = model(last_sequence).item()
                category_predictions.append(pred)
                last_sequence = torch.cat((last_sequence[:, 1:, :], torch.FloatTensor([[pred]]).view(1, 1, 1)), dim=1)

        # Save predictions
        predictions[category] = [{"date": date.strftime('%Y-%m-%d'), "sales": sales} for date, sales in zip(future_dates, category_predictions)]

    # Write predictions to JSON files
    os.makedirs(output_json_path, exist_ok=True)
    for category, preds in predictions.items():
        category_id = category.split('.')[0]  # Get category id from filename
        with open(os.path.join(output_json_path, f'json_category{category_id}.json'), 'w') as f:
            json.dump(preds, f, indent=4)

        # Plotting original and predicted data
        plt.figure(figsize=(10, 6))
        plt.scatter(future_dates, category_predictions, label='Predicted Sales', color='orange', s=10)
        plt.scatter(pd.date_range(start='2022-07-01', end='2023-06-30'), sales, label='Original Sales', color='blue', s=10)
        plt.xlabel('Date')
        plt.ylabel('Sales')
        plt.title(f'Sales Prediction for Category {category_id}')
        plt.legend()
        plt.savefig(os.path.join(graph_output_path, f'graph_category{category_id}.png'))
        plt.close()

# Main function to execute the workflow
if __name__ == "__main__":
    print("Starting LSTM Predictions...")
    input_dir = "../fujian/fujian2/groupByCategory"
    output_json_path = "./test/json_output"
    graph_output_path = "./test/graph_output"

    # Load data
    sales_data = load_data(input_dir)
    
    # Train LSTM and make predictions
    train_and_predict(sales_data, output_json_path, graph_output_path)
